import numpy as np
from scipy.signal import welch, csd
import scipy.signal
from scipy import signal, fftpack
import scipy.fft as fft
import math, zlib
from itertools import chain


# 速度积分：频域积分,更新时间：2024-11-29
def iomega(datain, dt, datain_type, dataout_type):
    """
    根据输入数据的类型和转换类型，对输入数据进行处理并返回转换后的数据。

    参数:
    - datain: 输入的数据数组。
    - dt: 时间步长，用于计算频率。
    - datain_type: 输入数据的类型，取值为1、2或3，表示不同的数据处理方式。
    - dataout_type: 输出数据的类型，取值为1、2或3，表示不同的数据处理方式。

    返回值:
    - dataout: 经过处理转换后的数据数组。

    异常:
    - 如果datain_type或dataout_type不在允许的范围内，将抛出ValueError。
    """
    # 检查输入数据类型和输出数据类型的有效性
    if datain_type < 1 or datain_type > 3:
        raise ValueError("Value for datain type must be 1, 2, or 3")
    elif dataout_type < 1 or dataout_type > 3:
        raise ValueError("Value for dataout type must be 1, 2, or 3")

    # 将输入数据转换为numpy数组，并计算所需的转换参数
    datain = np.array(datain)
    N = int(2 ** np.ceil(np.log2(len(datain))))
    df = 1 / (N * dt)
    Nyq = 1 / (2 * dt)

    # 计算虚数单位乘以角频率的数组
    iomega_array = 1j * 2 * np.pi * np.arange(-Nyq, Nyq, df)
    iomega_exp = dataout_type - datain_type

    # 对数据进行傅里叶变换，并进行必要的处理
    A = np.fft.fft(datain, N)
    A = np.fft.fftshift(A)
    for j in range(N):
        if iomega_array[j] != 0:
            A[j] = A[j] * (iomega_array[j] ** iomega_exp)
        else:
            A[j] = complex(0.0, 0.0)

    # 进行逆傅里叶变换，准备输出数据
    A = np.fft.ifftshift(A)
    dataout = np.fft.ifft(A)
    return dataout[: len(datain)]


# 时域转频域：傅里叶变换,更新时间：2024-11-29
# 傅里叶变换
def GetFFT(x, Fs, w_type="rectangular", s_type="Ewpk"):
    """
    Return the Fourier spectrum 'fftx' and spectrum power 'ffty' Parameters
    ---------- x: array_like, vibration data
    Fs: float, sampling frequency
    w_type: str-like, determining the window function ultilized, including 'rectangular','hanning', 'hamming','blackman', etc.
    s_type: 谱类型: 'Ewpk' 峰值, 'Eurms' 有效值
    """
    x = x - np.mean(x)
    x = befarr(x)
    p_l = 2.56
    L = np.size(x, 0)  # 数据长度
    N = np.power(2, np.ceil(np.log2(L)))  # 确定 FFT 长度
    N = int(N)
    xf = fft.fftfreq(N, 1 / Fs)[: int(N / p_l)]  # 确定 FFT 横坐标 /2.56，非 /2
    if w_type == "rectangular":  # 确定窗函数
        yf = fft.fft(x, N)
        yy = 2 * np.abs(yf[0 : int(N / p_l)] / L)  # 双边谱到单边谱
    elif w_type == "hanning":  # 确定窗函数
        w = np.hanning(L)
        wavedata_windowed = x * w
        yf = fft.fft(wavedata_windowed, N)  # 双边谱到单边谱
        yy = 2 * 2 * np.abs(yf[0 : int(N / p_l)] / L)
    elif w_type == "hamming":  # 确定窗函数
        w = np.hamming(L)
        wavedata_windowed = x * w
        yf = fft.fft(wavedata_windowed, N)  # 双边谱到单边谱
        yy = 1.852 * 2 * np.abs(yf[0 : int(N / p_l)] / L)
    elif w_type == "blackman":  # 确定窗函数
        w = np.blackman(L)
        wavedata_windowed = x * w
        yf = fft.fft(wavedata_windowed, N)
        yy = 2.381 * 2 * np.abs(yf[0 : int(N / p_l)] / L)  # 双边谱到单边谱
        yy[0] = 0
    if s_type == "Eurms":
        yy /= np.sqrt(2)
    # d = [[item1, item2] for item1, item2 in zip(xf, yy)]
    # return d  # 返回
    d = [[item1, round(item2, 4)] for item1, item2 in zip(xf, yy)]
    return d  # 返回


# 带通滤波,更新时间：2024-11-29
def GetFilter(x, Fs, fmin=0, fmax=0, type="bandpass"):
    """Return the filtered original signal.

    Parameters:
    -----------
    x : array_like
        Vibration data to filter.
    Fs : int or float
        Sampling frequency of vibration data.
    fmin : int or float
        Lower frequency for bandpass or highpass filters.
    fmax : int or float
        Upper frequency for bandpass or lowpass filters.
    type : {'bandpass', 'bandstop', 'lowpass', 'highpass'}
        Type of filter.

    Returns:
    --------
    x_filtered : array_like
        Filtered signal.
    """

    if fmin == fmax == 0:
        return x  # 如果没有指定频率范围，则返回原始信号

    if not (fmin == fmax == 0):
        if fmin <= 0:
            fmin = 0.00001  # 避免零频率导致的问题

    # 计算奈奎斯特频率
    nyquist_frequency = int(Fs / 2) - 1

    # 确保 fmax 不超过奈奎斯特频率
    if fmax > nyquist_frequency:
        fmax = nyquist_frequency

    # 根据滤波器类型设置Butterworth滤波器参数
    if type == "bandpass" or type == "bandstop":
        sos = signal.butter(4, [fmin, fmax], btype=type, fs=Fs, output="sos")
    elif type == "lowpass":
        sos = signal.butter(4, fmax, btype="lowpass", fs=Fs, output="sos")
    elif type == "highpass":
        sos = signal.butter(4, fmin, btype="highpass", fs=Fs, output="sos")

    # 应用滤波器
    x_filtered = signal.sosfilt(sos, x)

    if type == "bandstop":
        x_filtered = (
            x - x_filtered
        )  # 对于带阻滤波器，需要从原始信号中减去带通滤波后的结果

    return x_filtered


# 峰值去趋势
def detrend(z):
    # 计算信号的绝对值以找到峰值
    z_abs = np.abs(z)

    # 找到信号的峰值
    peaks, _ = scipy.signal.find_peaks(
        z_abs, height=np.max(z_abs) * 0.1
    )  # 高度阈值可以根据需要调整

    # 初始化一个列表来保存去趋势后的信号段
    detrended_segments = []

    # 对每一对连续的峰值之间的数据进行去趋势处理
    for i in range(len(peaks) - 1):
        start, end = peaks[i], peaks[i + 1]
        segment = z[start:end]

        # 对这段数据进行去趋势处理
        detrended_segment = scipy.signal.detrend(segment)

        # 添加到结果列表中
        detrended_segments.append(detrended_segment)

    # 如果最后一个峰值之后还有数据，则处理这部分数据
    if peaks[-1] < len(z):
        segment = z[peaks[-1] :]
        detrended_segment = scipy.signal.detrend(segment)
        detrended_segments.append(detrended_segment)

    # 合并所有去趋势后的段
    detrended_z = np.concatenate(detrended_segments)

    # 输出结果
    return detrended_z


def get_speed(x, fs, fmin=10, fmax=1000, filter_type="bandpass", resample_rate=2):
    # 重采样
    x = signal.decimate(x, resample_rate)
    fs = int(fs / resample_rate)
    # 去趋势处理
    x_detrended = signal.detrend(x)

    # 预处理滤波
    # if not (fmin == fmax == 0):
    #     if fmin<=0:
    #         fmin = 0.00001
    #     sos = signal.butter(4, [fmin, fmax], btype=filter_type, fs=fs, output='sos')
    #     x_detrended = signal.sosfilt(sos, x_detrended)
    x_detrended = GetFilter(x_detrended, fs, fmin, fmax, filter_type)

    # 积分操作
    velint = iomega(x_detrended, 1 / fs, 3, 2)

    # 再次去趋势处理
    velint_detrended = signal.detrend(velint)

    # 转换单位
    y_quzhiliu = velint_detrended * 1000

    return y_quzhiliu.real


# # 加速度转速度：速度积分,更新时间：2024-11-29
# def get_speed(x, fs, fmin=10, fmax=1000, filter_type="bandpass"):
#     """
#     返回经过滤波处理的原始信号。
#     参数：
#     ------
#     x: array_like
#        待进行傅里叶变换的振动数据。
#     fs: int or float
#         振动数据的采样频率，例如 51200。
#     fmin: int or float, 可选
#           停止频率和通过频率的下限，默认为10。
#     fmax: int or float, 可选
#           停止频率和通过频率的上限，默认为1000。
#     filter_type: str, 可选
#           指示滤波器的功能，默认为 'bandpass'，可选值包括 'bandpass', 'bandstop', 'lowpass', 'highpass'。

#     返回值：
#     -------
#     list
#         经过滤波处理的信号。
#     """
#     print("speed", x)
#     # 去趋势处理
#     x_detrended = signal.detrend(x)

#     # 预处理滤波
#     # if not (fmin == fmax == 0):
#     #     if fmin<=0:
#     #         fmin = 0.00001
#     #     sos = signal.butter(4, [fmin, fmax], btype=filter_type, fs=fs, output='sos')
#     #     x_detrended = signal.sosfilt(sos, x_detrended)
#     x_detrended = GetFilter(x_detrended, fs, fmin, fmax, filter_type)

#     # 积分操作
#     velint = iomega(x_detrended, 1 / fs, 3, 2)

#     # 再次去趋势处理
#     velint_detrended = signal.detrend(velint)

#     # 转换单位
#     y_quzhiliu = velint_detrended * 1000

#     return y_quzhiliu.real


def get_sudu(y_data, fs):
    y_data = befarr(y_data)
    max_index = y_data.index(max(y_data))
    left_arr = y_data[: max_index + 1]
    right_arr = y_data[max_index:]
    ret_left = [0]
    ret_right = [0]
    del_fs = 1 / fs
    l_right = len(right_arr)
    l_left = len(left_arr)
    print(max_index)
    if l_right >= 1:
        for i in range(l_right - 1):
            ret_right.append(
                ret_right[-1] + del_fs * (right_arr[i] + right_arr[i + 1]) / 2
            )
    if l_left >= 1:
        for i in range(l_left - 1):
            ret_left.append(
                ret_left[-1]
                - del_fs * (left_arr[l_left - i - 1] + left_arr[l_left - i - 2]) / 2
            )

    ret_left.reverse()
    ret = (np.round(np.array(ret_left[:-1] + ret_right) * 1000, 3)).tolist()
    l_ret = len(ret)
    ll = l_ret // 100
    for i in range(ll):
        ret[i * 100 : (i + 1) * 100] = ret[i * 100 : (i + 1) * 100] - np.mean(
            ret[i * 100 : (i + 1) * 100]
        )
    ret[ll * 100 :] = ret[ll * 100 :] - np.mean(ret[ll * 100 :])
    return ret


# 倒谱,更新时间：2024-11-29
def get_rep(x):
    # 计算频谱
    spectrum = np.fft.rfft(x)

    # 计算对数幅度谱
    log_spectrum = np.log(np.abs(spectrum) + 1e-21)

    # 计算倒谱
    cepstrum = np.fft.irfft(log_spectrum).real

    # 可选：滤除低频和高频部分
    cepstrum[0] = 0
    # cepstrum[-4:] = 0

    return cepstrum.tolist()[: int(len(cepstrum) / 2)]


# 包络谱,更新时间：2024-11-29
def GetEnvelope(
    x, Fs, low=0, high=0, filter_type="bandpass", w_type="hanning", s_type="Ewpk"
):
    # 去趋势处理
    x = signal.detrend(x)
    x = GetFilter(x, Fs, low, high, filter_type)
    ht = fftpack.hilbert(x)
    at = np.sqrt(np.square(ht) + np.square(x))
    at = at - np.mean(at)
    envelope = GetFFT(at, Fs, w_type, s_type)
    # envelope = [[item1, item2] for item1, item2 in zip(alpha, xSES)]
    return envelope


# 解压
def get_wave(wave):
    if type(wave) is list:
        return wave
    # 将解压缩后的字节数据转换回浮点数组
    restored_data = np.frombuffer(zlib.decompress(wave), dtype=np.float32).tolist()
    return [round(num, 4) for num in restored_data]


def befarr(arr):
    arr = np.asarray(arr)  # 确保输入是 NumPy 数组
    # print("arr",type(arr))
    length = arr.shape[0]  # 获取数组的长度

    if length & (length - 1) == 0:
        # 长度已经是 2 的 n 次幂，不需要补齐
        return arr
    else:
        # 计算补齐后的新长度
        padded_length = 1 << (length - 1).bit_length()
        # 使用 NumPy 的 pad 函数来补齐数组
        return np.pad(arr, (0, padded_length - length), mode="constant")


def get_csd(signal1, signal2, fs):
    """
    计算两个信号的互功率谱。

    参数:
    signal1, signal2 - 输入信号，一维numpy数组。
    fs - 采样频率
    N - 数组长度
    返回值:
    freqs - 频率数组
    cps - 互功率谱
    """
    min_length = min(len(signal1), len(signal2))
    signal1 = signal1[:min_length]
    signal2 = signal2[:min_length]
    freqs, cps = csd(signal1, signal2, fs=fs, nperseg=min_length)
    amplitude_spectrum = np.abs(cps)
    phase_spectrum = np.degrees(np.angle(cps))
    amplitude_spectrum = [
        [item1, round(item2, 4)] for item1, item2 in zip(freqs, amplitude_spectrum)
    ]
    phase_spectrum = [
        [item1, round(item2, 4)] for item1, item2 in zip(freqs, phase_spectrum)
    ]
    return amplitude_spectrum, phase_spectrum


def get_histogram(waveform, freq):
    energies = get_erenges(waveform, freq)
    # 提取能量值
    en = np.array([e[1] for e in energies])

    # 计算直方图
    max_en = np.max(en)
    bins = np.linspace(0, max_en, len(en) + 1)
    hist, _ = np.histogram(en, bins=bins)

    # 构造输出数据
    bin_centers = (bins[:-1] + bins[1:]) / 2
    bin_centers = bin_centers.astype(float)  # 将 bin_centers 转换为 float 类型
    data = [
        [float(bin_centers[i]), round(int(hist[i]), 4)] for i in range(len(hist))
    ]  # 将 bin_centers 转换为 float 类型
    return data


def get_erenges(waveform, freq):
    acc = np.array(waveform)
    # 计算能量
    en = np.square(acc)
    # 构造输出数据
    times = np.arange(len(waveform)) / freq
    data = [[times[i], round(en[i], 4)] for i in range(len(waveform))]
    return data


def bef_get_wave(data, freq):
    time_step = 1000 / freq
    return [[round(time_step * i, 3), e] for i, e in enumerate(data)]


# 计算波形特征值
def get_tzz(arr, tag):
    ret = {}
    mean = round(np.mean(arr), 4)
    ret.update({tag + "_mean": mean})
    max = round(np.max(arr), 4)
    ret.update({tag + "_max": max})
    min = round(np.min(arr), 4)
    ret.update({tag + "_min": min})
    p_p = round(max - min, 4)
    ret.update({tag + "_pp": p_p})
    rms = round(np.sqrt(np.sum(np.square(arr)) / len(arr)), 4)
    ret.update({tag + "_rms": rms})
    peak = round(np.max(np.abs(arr)), 4)
    ret.update({tag + "_peak": peak})
    root = round(math.pow(np.sum(np.sqrt(np.abs(arr))) / len(arr), 2), 4)
    ret.update({tag + "_root": root})
    ave = round(np.sum(np.abs(arr)) / len(arr), 4)
    ret.update({tag + "_ave": ave})
    sha = round(rms / ave, 4)
    ret.update({tag + "_sha": sha})
    cre = round(peak / rms)
    ret.update({tag + "_cre": cre})
    kur = round(np.sum(np.power(arr, 4)) / (len(arr) * pow(rms, 4)), 4)
    ret.update({tag + "_kur": kur})
    ske = round(np.sum(np.power(arr, 3)) / (len(arr) * pow(rms, 3)), 4)
    ret.update({tag + "_ske": ske})
    imp = round(peak / ave, 4)
    ret.update({tag + "_imp": imp})
    cle = round(peak / root, 4)
    ret.update({tag + "_cle": cle})
    return ret
